import cv2
import os

if __name__ == '__main__':
    # 视频路径
    video_path = '/home/lidar/foursteps/Foursteps/yolov8/CreateDataset/record/video/rods/video1.mp4'
    # 输出目录
    output_dir = '/home/lidar/foursteps/Foursteps/yolov8/CreateDataset/dataset/YOLODataset/rod-image'
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print("无法打开视频文件")
        exit()
    
    # 获取视频的帧率（FPS）
    fps = cap.get(cv2.CAP_PROP_FPS)
    print(f"视频帧率: {fps} FPS")
    
    # 设置保存间隔（每20帧保存一次）
    frame_interval = 10
    frame_count = 0
    saved_count = 0
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break  # 视频结束
        
        # 按间隔保存帧
        if frame_count % frame_interval == 0:
            output_path = os.path.join(output_dir, f"{saved_count}.png")
            cv2.imwrite(output_path, frame)
            saved_count += 1
        
        frame_count += 1
    
    # 释放资源
    cap.release()
    print(f"处理完成！共处理 {frame_count} 帧，保存了 {saved_count} 张图像。")